import faiss

import numpy as np
d = 64       # dimension
nb = 100000  # database size
nq = 10000   # number of queries
np.random.seed(1234)

xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000

# print(xb[0:10])
# print(xq[0:10])

index = faiss.IndexFlatL2(d)  # build the index

print(index.is_trained)
index.add(xb)                 # add vectors to the index
print(index.ntotal)

k = 4   # search k nearest neighbors
D, I = index.search(xq, k)
print(I[:5])   # neighbors of the 5 first queries
print(D[:5])  # neighbors of the 5 first queries

n_list = 100
k = 4
m = 8 # PQ compress encoding length to be 8 bytes
quantizer = faiss.IndexFlatL2(d)  # another index
# index = faiss.IndexIVFFlat(quantizer, d, n_list, faiss.METRIC_L2)
# index = faiss.IndexIVFPQ(quantizer, d, n_list, m, 8)  # 8 specifies that each sub-vector is encoded as 8 bits
index = faiss.index_factory(d, "IVF100,PQ8", faiss.METRIC_L2)
assert(not index.is_trained)

index.train(xb)

assert index.is_trained

index.add(xb)
D, I = index.search(xq, k)
print(I[:5])
index.nprobe = 10
D, I = index.search(xq, k)
print(I[:5])




